#!/usr/bin/env python
# coding: utf-8

"""
work_conf.py
create by cuiyu 2025-03-27
说明：
 * 配置文件读写

用法:
  work_conf.py

更新:
2025-03-27:
"""

from pathlib import Path
from dataclasses import dataclass, asdict
from ruamel.yaml import YAML

@dataclass
class CPipelineConfig:
    """
    训练工作流水线工作参数
    """
    yaml_tag = u'!Pipe'  # 定义 YAML 标签
    datainit: bool = True  # True = 重新生成标注文件 (将 labelme 标注文件转换为 yolo 格式的 txt 文件)
    conf: float = 0.4
    img_root_train: str = "/aaby/netdisk/mydata/ai-data/dataset/priv/井盖病害/训练数据/img01"
    data_root_predict: str = "/aaby/netdisk/mydata/ai-data/dataset/priv/井盖病害/自动标注数据"


@dataclass
class CModelConfig:
    """
    训练模型参数，默认为 debug 参数配置
    """
    model_file: str = "/aaby/netdisk/mydata/ai-data/model/ultralytics/od/yolo11n.pt"
    # train_file: str = ""
    batch: int = 2
    epochs: int = 2
    patience: int = 0       # 早停参数
    workers: int = 0
    freeze: int = 10   
    imgsz: int = 640
    rect: bool = True
    profile: bool = True      # 训练期间对 ONNX 和 TensorRT 的速度进行分析，有助于优化模型部署。
    optimizer: str = "SGD"    # 优化器，默认使用 SGD，Adam 效果一般
    lr0: float = 0.001        # 初始学习率（实测道路缺陷最佳值）; 初始学习率设为0.001（低于默认0.01），避免大尺寸图像训练初期梯度爆炸
    lrf: float = 0.1          # 最终学习率（实测道路缺陷最佳值）; 最终学习率设为0.01（低于默认0.1），避免训练后期学习率过快
    momentum: float = 0.937   # SGD 优化器的动量参数，用于加速梯度下降过程。

class CAppConfig:
    """
    训练工作流程配置类
    """

    def __init__(self):
        # 创建 YAML 对象
        self.yaml = YAML(typ="safe")
        self.yaml.preserve_quotes = True  # 保留字符串引号（可选）
        self.yaml.indent(mapping=2, sequence=4, offset=2)  # 设置缩进格式
        self.yaml.allow_unicode = True  # 关键：禁用 Unicode 转义

        # 配置参数
        self.model_cfg: CModelConfig = CModelConfig()
        self.pipe_cfg: CPipelineConfig = CPipelineConfig()
        self.yaml_file: Path = Path("")


    def load(self, yaml_file: Path) -> None:
        """
        加载配置文件，将配置文件中的内容加载到全局变量中，方便后续使用。
        配置文件的格式为："
        """
        self.yaml_file = yaml_file
            

        # 检查配置文件是否存在，如果不存在则使用默认配置
        if not yaml_file.exists():
            print(f"配置文件 {yaml_file} 不存在，使用默认配置")
            self.save(yaml_file)
            return

        # 读取配置文件
        with open(yaml_file, "r", encoding="utf-8") as file:
            # 使用safe_load()来解析YAML文件内容到Python字典
            yaml_data = self.yaml.load(file)

            # 1. pipe_conf 直接解包，无法处理嵌套的 dataclass 或复杂类型（如列表、Path、其他dataclass实例）
            self.pipe_cfg = CPipelineConfig(**yaml_data["pipe"]) 

            # 2. train_conf，直接解包
            self.model_cfg = CModelConfig(**yaml_data["model"])

        pass

    def save(self, yaml_file: Path) -> None:
        """
        保存配置文件，将全局变量中的内容保存到配置文件中。
        """
        yaml_data: dict = {}
        # 1. pipe_cfg
        yaml_data["pipe"]= asdict(self.pipe_cfg)

        # 2. model_cfg
        yaml_data["model"]= asdict(self.model_cfg)

        # print(yaml_data)
        with open(yaml_file, "w", encoding="utf-8") as f:
            # YAML 对象的配置需要在每次 dump 前重新设置
            yaml = YAML(typ="rt")        # 一定要用 rt，不能用 safe，才能保证格式化输出正确
            yaml.preserve_quotes = True  # 保留字符串引号（可选）
            yaml.indent(mapping=2, sequence=4, offset=2)  # 设置缩进格式
            yaml.allow_unicode = True  # 关键：True 禁用 Unicode 转义，确保中文等 unicode 直接输出而非转义输出
            yaml.dump(yaml_data, f)


APP_CFG: CAppConfig = CAppConfig()


def _test():
    import sys
    from rich import print
    # 全局变量
    _APP_FILE = Path(sys.argv[0])
    _APP_ROOT = _APP_FILE.parent

    conf_file = _APP_ROOT / "app-debug.yaml"
    # APP_CONF.save(conf_file)
    APP_CFG.load(conf_file)
    print(APP_CFG.pipe_cfg)
    print("="*20)
    print(APP_CFG.model_cfg)
    # print("="*20)
    # print(APP_CONF.pipe_conf.img_root_train.parent)

    pass

if __name__ == "__main__":
    _test()